import math
from collections import Counter
from operator import itemgetter


class KNN:
    def __init__(self, k=3, task='classification'):
        self.k = k
        self.task = task
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):
        self.X_train = X
        self.y_train = y

    def _euclidean_distance(self, a, b):
        #计算两个向量之间的欧氏距离
        return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

    def _get_k_neighbors(self, x):
        
        #获取距离目标样本最近的 k 个邻居

        
        distances = []

        #与所有训练样本计算距离
        for i, train_sample in enumerate(self.X_train):
            dist = self._euclidean_distance(x, train_sample)
            distances.append((i, dist))

        #排序取最近的 k 个
        distances.sort(key=itemgetter(1))
        neighbors = distances[:self.k]

        return neighbors

    def predict(self, X):
        #预测新样本标签（分类或回归）
        predictions = []

        for sample in X:
            neighbors = self._get_k_neighbors(sample)

            if self.task == 'classification':
                # 多数投票
                neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
                most_common = Counter(neighbor_labels).most_common(1)[0][0]
                predictions.append(most_common)

            elif self.task == 'regression':
                # 平均值
                neighbor_values = [self.y_train[idx] for idx, _ in neighbors]
                predictions.append(sum(neighbor_values) / len(neighbor_values))

        return predictions

    def predict_proba(self, X):
        #返回每个样本属于各类别的概率（仅分类任务）
        if self.task != 'classification':
            raise ValueError("predict_proba 仅适用于分类任务")

        probabilities = []

        for sample in X:
            neighbors = self._get_k_neighbors(sample)
            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]

            total = len(neighbors)
            counts = Counter(neighbor_labels)

            # 构造类别概率字典
            prob_dict = {}
            classes = set(self.y_train)
            for c in classes:
                prob_dict[c] = counts.get(c, 0) / total

            probabilities.append(prob_dict)

        return probabilities
